package com.aiwiown.snackmq.broker.handler;

import com.aiwiown.snackmq.broker.auth.AuthorizationService;
import com.aiwiown.snackmq.common.auth.Action;
import com.aiwiown.snackmq.common.exception.AuthorizationException;
import com.aiwiown.snackmq.common.handler.MessageHandler;
import com.aiwiown.snackmq.common.handler.RequestContext;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import com.aiwiown.snackmq.common.util.MessageConverter;
import com.aiwiown.snackmq.storage.exception.PartitionNotFoundException;
import com.aiwiown.snackmq.storage.partition.LogPartition;
import com.aiwiown.snackmq.storage.service.PartitionedStorageService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 【新增】处理消费者拉取消息请求的 Handler。
 * 增加了健壮的异常处理，确保任何失败都会向客户端返回明确的错误信息。
 */
@Slf4j
@RequiredArgsConstructor // 【修复】使用 Lombok 简化构造函数
public class PullMessageHandler implements MessageHandler {

    private final PartitionedStorageService storageService;
    private final AuthorizationService authorizationService; // 【修复】增加授权服务依赖

    @Override
    public MessageType getMessageType() {
        return MessageType.PULL_REQUEST;
    }

    @Override
    public void handle(RequestContext ctx, Message request) {
        try {
            String topic = request.getTopic();
            String consumerGroup = request.getProperty("consumerGroup");
            // 1. 【修复】增加授权检查，包括消费者组权限
            authorizationService.authorize(
                    request.getProperty("authToken"),
                    Action.CONSUME,
                    new AuthorizationService.Resource("TOPIC", topic),
                    new AuthorizationService.Resource("CONSUMER_GROUP", consumerGroup)
            );

            // 2. 检查请求参数的有效性
            if (topic == null || request.getProperty("partition") == null || request.getProperty("offset") == null || request.getProperty("maxMessages") == null) {
                throw new IllegalArgumentException("Pull request is missing required properties (topic, partition, offset, maxMessages).");
            }

            int partitionId = Integer.parseInt(request.getProperty("partition"));
            long offset = Long.parseLong(request.getProperty("offset"));
            int maxMessages = Integer.parseInt(request.getProperty("maxMessages"));

            TopicPartition tp = new TopicPartition(topic, partitionId);

            // 3. 检查分区是否存在
            LogPartition partition = storageService.getLogPartition(tp);
            if (partition == null) {
                throw new PartitionNotFoundException("Partition " + tp + " is not active on this broker.");
            }

            // 4. 从存储中读取消息
            List<ByteBuffer> messageBuffers = partition.readBatchZeroCopy(offset, maxMessages);

            // 5. 构建成功或空响应
            Message.MessageBuilder responseBuilder = Message.builder()
                    .messageId(request.getMessageId())
                    .type(MessageType.RESPONSE);

            if (messageBuffers.isEmpty()) {
                responseBuilder.status(MessageStatus.EMPTY);
            } else {
                List<Message> messages = messageBuffers.stream()
                        .map(MessageConverter::fromByteBuffer)
                        .collect(Collectors.toList());
                responseBuilder.status(MessageStatus.SUCCESS).messages(messages);
            }
            ctx.writeResponse(responseBuilder.build());

        } catch (AuthorizationException e) { // 【修复】捕获并处理授权异常
            log.warn("Authorization failed for PULL request: {}", e.getMessage());
            Message errorResponse = Message.builder()
                    .messageId(request.getMessageId())
                    .type(MessageType.RESPONSE)
                    .status(MessageStatus.FAILED)
                    .errorMessage("Permission denied: " + e.getMessage())
                    .build();
            ctx.writeResponse(errorResponse);
        } catch (IllegalArgumentException | PartitionNotFoundException e) {
            log.warn("Failed to process PULL request due to invalid argument or partition not found: {}", e.getMessage());
            Message errorResponse = Message.builder()
                    .messageId(request.getMessageId())
                    .type(MessageType.RESPONSE)
                    .status(MessageStatus.FAILED)
                    .errorMessage(e.getMessage())
                    .build();
            ctx.writeResponse(errorResponse);
        } catch (IOException e) {
            log.error("Failed to read from storage for PULL request: {}", request, e);
            Message errorResponse = Message.builder()
                    .messageId(request.getMessageId())
                    .type(MessageType.RESPONSE)
                    .status(MessageStatus.FAILED)
                    .errorMessage("Broker storage error while reading messages.")
                    .build();
            ctx.writeResponse(errorResponse);
        } catch (Exception e) {
            log.error("Unexpected error while processing PULL request: {}", request, e);
            Message errorResponse = Message.builder()
                    .messageId(request.getMessageId())
                    .type(MessageType.RESPONSE)
                    .status(MessageStatus.FAILED)
                    .errorMessage("Unexpected broker error: " + e.getMessage())
                    .build();
            ctx.writeResponse(errorResponse);
        }
    }
}